import grpc
from p4.v1 import p4runtime_pb2
from p4.v1 import p4runtime_pb2_grpc
from google.protobuf import text_format
import argparse
import socket
import struct

def ip_to_bytes(ip):
    """将IPv4地址转换为字节表示"""
    return socket.inet_aton(ip)

def mac_to_bytes(mac):
    """将MAC地址转换为字节表示"""
    return bytes.fromhex(mac.replace(':', ''))

def add_acl_rule(stub, p4info, device_id, src_ip, dst_ip, protocol, action):
    """添加ACL规则到交换机"""
    # 创建写入请求
    request = p4runtime_pb2.WriteRequest()
    request.device_id = device_id
    update = request.updates.add()
    update.type = p4runtime_pb2.Update.INSERT
    
    # 获取表ID和动作ID
    table_id = p4info.tables_by_name['acl_table'].id
    
    if action == 'allow':
        action_id = p4info.actions_by_name['allow'].id
    else:  # drop
        action_id = p4info.actions_by_name['drop'].id
    
    # 构建表项
    table_entry = update.entity.table_entry
    table_entry.table_id = table_id
    
    # 添加匹配字段
    for field_name, field_value in [
        ('hdr.ipv4.srcAddr', ip_to_bytes(src_ip)),
        ('hdr.ipv4.dstAddr', ip_to_bytes(dst_ip)),
        ('hdr.ipv4.protocol', struct.pack('!B', protocol))
    ]:
        match = table_entry.match.add()
        field_descriptor = p4info.tables_by_name['acl_table'].match_fields_by_name[field_name]
        match.field_id = field_descriptor.id
        match.exact.value = field_value
    
    # 设置动作
    action = table_entry.action.action
    action.action_id = action_id
    
    # 发送请求
    stub.Write(request)
    print(f"已添加ACL规则: {src_ip} -> {dst_ip} (协议: {protocol}) [{action}]")

def main():
    parser = argparse.ArgumentParser(description='使用P4Runtime下发ACL规则')
    parser.add_argument('--grpc-addr', type=str, default='localhost:50051',
                        help='P4Runtime服务器地址')
    parser.add_argument('--device-id', type=int, default=1,
                        help='设备ID')
    parser.add_argument('--p4info', type=str, required=True,
                        help='P4Info文件路径')
    parser.add_argument('--rule-file', type=str, required=True,
                        help='ACL规则文件路径')
    
    args = parser.parse_args()
    
    # 连接到P4Runtime服务器
    channel = grpc.insecure_channel(args.grpc_addr)
    stub = p4runtime_pb2_grpc.P4RuntimeStub(channel)
    
    # 读取P4Info文件
    with open(args.p4info, 'r') as f:
        p4info = p4runtime_pb2.P4Info()
        text_format.Parse(f.read(), p4info)
    
    # 读取ACL规则文件
    with open(args.rule_file, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
                
            parts = line.split()
            if len(parts) != 4:
                print(f"跳过无效规则: {line}")
                continue
                
            src_ip, dst_ip, protocol, action = parts
            try:
                protocol = int(protocol)
                add_acl_rule(stub, p4info, args.device_id, src_ip, dst_ip, protocol, action)
            except Exception as e:
                print(f"添加规则失败: {line} - {str(e)}")

if __name__ == "__main__":
    main()